#include <mpi.h>
#include <stdio.h>
#include <stdlib.h>
#include "utils.h"


// 本例演示MPI通信子相关的API，主要包括:
//
// MPI_Comm_group, MPI_Comm_create, MPI_Comm_free;
// MPI_Comm_dup, MPI_Comm_split, MPI_Comm_compare；
// MPI_Group_excl, MPI_Group_incl, MPI_Group_free;
//  
//
int main(int argc, char** argv) {
    int rank, size, scount1, scount2, comp_res;
    MPI_Init(&argc, &argv);

    scount1 = 5;
    int send_buf1[scount1] = {1, 2, 3, 4, 5};
    int recv_buf1[scount1] = {0};

    scount2 = 6;
    int send_buf2[scount2] = {6, 7, 8, 9, 10, 11};
    int recv_buf2[scount2] = {0};

    MPI_Group grp_world, grp;
    MPI_Comm  comm_src, comm_cpy;
    MPI_Comm  comm_sp;
    static int ranks[1] = {0};                  // 创建新组时被排除的进程列表

    MPI_Status status;
    // MPI_Comm_rank(MPI_COMM_WORLD, &rank);    // 全局通信子下的进程编号
    // MPI_Comm_size(MPI_COMM_WORLD, &size);    // 全局通信子下的进程数量

    MPI_Comm_group(MPI_COMM_WORLD, &grp_world);
    MPI_Group_rank(grp_world, &rank);          // 与MPI_Comm_rank(MPI_COMM_WORLD, &rank)等价
    MPI_Group_size(grp_world, &size);          // 与MPI_Comm_size(MPI_COMM_WORLD, &size)等价

    INFO(" >> rank: %d, size: %d;\n", rank, size);

    //! 1. 创建一个不包含进程0的新的进程组
    //    grp_world: 原始进程组 
    //    1:         要排除的进程数量
    //    ranks:     要排除的进程号 
    //    &grp:      新的进程组
    MPI_Group_excl(grp_world, 1, ranks, &grp);

    //! 2. 基于新的进程组创建一个通信子
    MPI_Comm_create(MPI_COMM_WORLD, grp, &comm_src);      // comm_src 穿件的新的通信子

    if (rank != 0) {
        MPI_Comm_dup(comm_src, &comm_cpy);                // 复制通信子
        MPI_Comm_compare(comm_src, comm_cpy, &comp_res);  // 比较两个通信子
        INFO(">> rank: %d, compare result: %d;\n", rank, comp_res); // 1,2,3 输出为1
    }


    // rank%2相同的进程会被分入统一进程中
    // rank 根据原始全局通信子下的进程编号大小在新的通信子下排序
    // TODO: comm_sp 包含了两个通信子
    MPI_Comm_split(MPI_COMM_WORLD, rank%2, rank, &comm_sp);
    int newrank, newsize;
    MPI_Comm_rank(comm_sp, &newrank);    // 全局通信子下的进程编号
    MPI_Comm_size(comm_sp, &newsize);    // 全局通信子下的进程数量
    INFO(">> rank: %d, color:%d, newrank: %d, newsize: %d;\n", rank, rank%2, newrank, newsize);


    //! MPI_Reduce的结果只会存到某一个进程, 这里的根进程时rank=0
    // [INFO] Reduce Sum >> rank: 3, [ 0,  0,  0,  0,  0,  0];
    // [INFO] Reduce Sum >> rank: 2, [ 0,  0,  0,  0,  0,  0];
    // [INFO] Reduce Sum >> rank: 1, [ 0,  0,  0,  0,  0,  0];
    // [INFO] Reduce Sum >> rank: 0, [24, 28, 32, 36, 40, 44];
    MPI_Reduce(send_buf2, recv_buf2, scount2, MPI_INT, MPI_SUM, 0, MPI_COMM_WORLD);
    INFO("Reduce Sum >> rank: %d, [%2d, %2d, %2d, %2d, %2d, %2d];\n", rank, 
        recv_buf2[0], recv_buf2[1], recv_buf2[2], recv_buf2[3],recv_buf2[4],recv_buf2[5]);

    MPI_Group_free(&grp);
    MPI_Group_free(&grp_world);
    MPI_Comm_free(&comm_sp);

    MPI_Finalize();
    return 0;
}